mod context;
use riscv::register::{
    sie,
    stval,
    stvec,
    scause::{
        self,
        Trap,
        Exception,
        Interrupt,
    },
    mtvec::TrapMode,
};
use crate::task::{
    exit_current_and_run_next,
    suspend_current_and_run_next,
    current_user_token,
    current_trap_ctx,
};

use crate::config::{TRAP_CONTEXT, TRAMPOLINE};

pub use context::TrapContext;

use crate::syscall::syscall; 

global_asm!(include_str!("trap.S"));

pub fn init() {
    set_kernel_trap_entry();
}

fn set_kernel_trap_entry() {
    unsafe {
        stvec::write(trap_from_kernel as usize, TrapMode::Direct);
    }
}

fn set_user_trap_entry() {
    unsafe {
        stvec::write(TRAMPOLINE as usize, TrapMode::Direct);
    }
}

// 设置了 sie.stie 使得 S 特权级时钟中断不会被屏蔽
pub fn enable_timer_interrupt() {
    unsafe { sie::set_stimer(); }
}

#[no_mangle]
pub fn trap_handler() -> ! {
    // 陷入 Trap 之后, 由硬件完成 sstatus sepc scause, stval 等的填写
    set_kernel_trap_entry();
    let ctx = current_trap_ctx();
    let scause = scause::read(); // 描述 Trap 的原因
    let stval = stval::read(); // 给出 Trap 附加信息
    match scause.cause() {
        Trap::Exception(Exception::UserEnvCall) => {
            ctx.sepc += 4; // ecall: 4字节长; 移动 4 个字节, 使其指向 ecall 下一条指令
            ctx.x[10] = syscall(ctx.x[17], [ctx.x[10], ctx.x[11], ctx.x[12]]) as usize;
        }
        Trap::Exception(Exception::StoreFault) |
        Trap::Exception(Exception::StorePageFault) |
        Trap::Exception(Exception::InstructionPageFault) |
        Trap::Exception(Exception::LoadFault) | 
        Trap::Exception(Exception::LoadPageFault) => {
            error!("[Kernel] PageFault in application, core dumped.");
            exit_current_and_run_next();
        }
        Trap::Exception(Exception::IllegalInstruction) => {
            error!("[Kernel] IllegalInstruction in application, core dumped.");
            debug!("id = {} args = [{}, {}, {}]", ctx.x[17], ctx.x[10], ctx.x[11], ctx.x[12]);
            exit_current_and_run_next();
        }
        Trap::Interrupt(Interrupt::SupervisorTimer) => {
            suspend_current_and_run_next();
        }
        _ => {
            panic!("Unsupported trap {:?}, stval = {:#x}", scause.cause(), stval);
        }
    }
    trap_return();
}

#[no_mangle]
pub fn trap_return() -> ! {
    set_user_trap_entry();
    let trap_cx_ptr = TRAP_CONTEXT;
    let user_satp = current_user_token();
    extern "C" {
        fn __alltraps();
        fn __restore();
    }
    let restore_va = __restore as usize - __alltraps as usize + TRAMPOLINE;
    unsafe {
        llvm_asm!("fence.i" :::: "volatile");
        llvm_asm!("jr $0" :: "r"(restore_va), "{a0}"(trap_cx_ptr), "{a1}"(user_satp) :: "volatile");
    }
    panic!("Unreachable in back_to_user!");
}

#[no_mangle]
pub fn trap_from_kernel() -> ! {
    panic!("a trap from kernel!");
}